import asyncio, json
import websockets
import struct
import opuslib
from pydub import AudioSegment
import numpy as np
from typing import Union, Dict, Any

"""
ffmpeg -i ./wav/gun.mp3 -ar 16000 -ac 1 -f s16le gun.pcm

"""

class NewBinaryProtocol:
  FORMAT_HEADER = '>BBH'  # Big-endian byte order, unsigned char, unsigned char, unsigned short
  HEADER_SIZE = struct.calcsize(FORMAT_HEADER)

  def __init__(self, type: int, reserved: int, payload: Union[bytes, str]):
    self.type = type
    self.reserved = reserved
    if isinstance(payload, str):
      self.payload = payload.encode('utf-8')
    else:
      self.payload = payload
    self.payload_size = len(self.payload)

  def serialize(self) -> bytes:
    header = struct.pack(self.FORMAT_HEADER, self.type, self.reserved, self.payload_size)
    return header + self.payload

  @classmethod
  def deserialize(cls, data: bytes) -> 'BinaryProtocol':
    header = data[:cls.HEADER_SIZE]
    payload = data[cls.HEADER_SIZE:]
    type, reserved, payload_size = struct.unpack(cls.FORMAT_HEADER, header)
    if len(payload) != payload_size:
      raise ValueError("Payload size does not match the expected length")
    return cls(type, reserved, payload)

  def to_dict(self) -> Dict[str, Any]:
    return {
      'type': self.type,
      'reserved': self.reserved,
      'payload_size': self.payload_size,
      'payload': self.payload.decode('utf-8') if isinstance(self.payload, bytes) else self.payload
    }

  @classmethod
  def from_dict(cls, d: Dict[str, Any]) -> 'BinaryProtocol':
    return cls(d['type'], d['reserved'], d['payload'])

# 定义采样率和声道数
sample_rate = 16000
channels = 1
frame_size = int(sample_rate * 60 / 1000)  # 每帧 60 毫秒

# 创建 Opus 编码器
encoder = opuslib.Encoder(sample_rate, channels, opuslib.APPLICATION_AUDIO)

# 读取 PCM 文件
def read_pcm_file(file_path):
  with open(file_path, 'rb') as f:
    pcm_data = f.read()
  return np.frombuffer(pcm_data, dtype=np.int16)

# 将 PCM 数据分割成帧
def split_into_frames(pcm_data, frame_size):
  return [pcm_data[i:i + frame_size] for i in range(0, len(pcm_data), frame_size)]


class BinaryProtocol:
  def __init__(self, data):
    self.data = data
    self.type = None
    self.reserved = None
    self.payload_size = None
    self.payload = None

    self.unpack()

  def unpack(self):
    # 解包头部信息
    header = self.data[:4]
    # self.type, self.reserved, self.payload_size = struct.unpack('=BB>H', header)
    self.type, self.reserved = struct.unpack('=BB', header[:2])
    self.payload_size, = struct.unpack('>H', header[2:])
    
    # 根据payload_size取出payload
    self.payload = self.data[4:4+self.payload_size]

  def to_dict(self):
    return {
      'type': self.type,
      'reserved': self.reserved,
      'payload_size': self.payload_size,
      # 'payload': self.payload
    }

  def __str__(self):
    return f"Type: {self.type}, Reserved: {self.reserved}, Payload Size: {self.payload_size}, Payload: {self.payload}"


def is_binary(data):
  return isinstance(data, bytes)

# 远程WebSocket服务端的URL
REMOTE_WS_URL = 'wss://api.tenclass.net/xiaozhi/v1/'

# 存储所有活动连接
connected_clients = set()

# 连接到远程WebSocket服务端的实例
remote_ws = None

extra_headers={
  "Authorization": "Bearer test-token",
  "Device-Id": "f0:9e:9e:0e:ad:b0",
  "Protocol-Version": 3
}

async def connect_to_remote():
  """连接到远程WebSocket服务端"""
  global remote_ws
  remote_ws = await websockets.connect(REMOTE_WS_URL, extra_headers=extra_headers)
  print("Connected to remote WebSocket server")

async def forward_to_remote(message):
  """将消息转发到远程WebSocket服务端"""
  if remote_ws and remote_ws.open:
    await remote_ws.send(message)
    if is_binary(message):
      protocol = BinaryProtocol(message)
      print("Message forwarded to remote server:: ", protocol.to_dict())
    else:
      print(f"Message forwarded to remote server:: {message}")

async def send_custom_msg():
  # 读取 PCM 文件
  pcm_file_path = './wav/gun.pcm'
  pcm_data = read_pcm_file(pcm_file_path)

  # 将 PCM 数据分割成帧
  frames = split_into_frames(pcm_data, frame_size)

  for frame in frames:
    # 编码 PCM 帧为 Opus 数据
    opus_frame = encoder.encode(frame.tobytes(), frame_size)

    # 序列化 BinaryProtocol
    audio_message = NewBinaryProtocol(type=0, reserved=0, payload=opus_frame)
    audio_serialized = audio_message.serialize()

    # 发送 Opus 数据
    await broadcast_to_clients(audio_serialized)
    print(f"Sent frame of size {len(opus_frame)} bytes")



intercept = False
async def forward_from_remote_to_clients():
  """从远程转发到客户端"""
  global intercept
  """从远程WebSocket服务端接收消息并转发给所有客户端"""
  while True:
    if remote_ws and remote_ws.open:
      try:
        message = await remote_ws.recv()
        if is_binary(message) == False and 'sentence' in message:
          msg = json.loads(message)
          print(msg)
          if msg['state'] == 'sentence_start':
            intercept = True
            msg['text'] = '今天天气不好!所以,滚蛋!'
            await broadcast_to_clients(json.dumps(msg, ensure_ascii=False))
            # 自定义音频发送
            await send_custom_msg()
            continue
          elif msg['state'] == 'sentence_end':
            msg['text'] = '今天天气不好!所以,滚蛋!'
            intercept = False
            await broadcast_to_clients(json.dumps(msg, ensure_ascii=False))
            continue
        if intercept == True:
          continue
        
        if is_binary(message):
          protocol = BinaryProtocol(message)
          print("Received message from remote server: ", protocol.to_dict())
        else:
          print(f"Received message from remote server: {message}")
        await broadcast_to_clients(message)
      except websockets.exceptions.ConnectionClosed:
        print("Remote server connection closed")
        break
    else:
      await asyncio.sleep(1)  # 每秒检查一次连接状态

async def broadcast_to_clients(message):
  """将消息广播给所有已连接的客户端"""
  if connected_clients:  # 如果有连接的客户端
    await asyncio.wait([client.send(message) for client in connected_clients])



async def handler(websocket, path):
  """处理客户端连接"""
  connected_clients.add(websocket)
  try:
    async for message in websocket:
      # print(f"Received message from client: {message}")
      if is_binary(message):
        protocol = BinaryProtocol(message)
        print("protocol Received message from client: ", protocol.to_dict())
        # opus_to_wav(protocol.to_dict()['payload'], f"./wav/{idx}.wav")
      else:
        print(f"Received message from client: {message}")
        
      # 转发消息到远程WebSocket服务端
      await forward_to_remote(message)
  except websockets.exceptions.ConnectionClosed as e:
    print("A client just disconnected.")
  finally:
    connected_clients.remove(websocket)

async def main():
  # 启动中转服务器
  start_server = websockets.serve(handler, "0.0.0.0", 18765)
  await start_server
  print("Server started on ws://0.0.0.0:18765")

  # 连接到远程WebSocket服务端
  await connect_to_remote()

  # 开始从远程服务端接收消息并转发给客户端
  await forward_from_remote_to_clients()

asyncio.run(main())